# -*- coding: utf-8 -*-
"""
 * @file main.py
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2024-05-31
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi. This file can not
 * be copied or distributed without the permission of Lynxi Technologies Co., Ltd.
"""

from collections import namedtuple
import multiprocessing
import signal
import sys
import json
import ipeParamYolov8

sys.path.append("../../../")
from ctypes import *
import argparse
import common.python.common as common
import struct
import pylynchipsdk as sdk
from common.python.plugin_utils import *
from common.python.infer_process import *
from common.python.callback_data_struct import *
import ctypes


# Create a shared boolean value to stop running thread
cancel_flag = multiprocessing.Value("b", False)
pythonapi.PyCapsule_GetPointer.restype = c_void_p
pythonapi.PyCapsule_GetPointer.argtypes = [py_object, c_char_p]

UINT8_PTR = ctypes.POINTER(ctypes.c_uint8)

def save_boxinfo_cb(params):
    output_path = params[0]
    boxex_info = params[1]
    eos = params[2]
    if eos:
        return 0

    dst_img_size = ctypes.sizeof(LynSegBoxesInfo)
    host_buf_arr = np.ones(dst_img_size, dtype=np.uint8)
    host_buf = sdk.lyn_numpy_to_ptr(host_buf_arr)
    ret = sdk.lyn_memcpy(
        host_buf, boxex_info, dst_img_size, sdk.lyn_memcpy_dir_t.ServerToClient
    )
    common.error_check(ret, "save_boxinfo_cb lyn_memcpy")
    host_buf_c = pythonapi.PyCapsule_GetPointer(host_buf, None)
    c_box = ctypes.cast(host_buf_c, ctypes.POINTER(LynSegBoxesInfo)).contents
    dump_box_json(output_path, c_box)
    return 0


def dump_box_json(output_path, c_box):
    box_dict = {"boxesnum": c_box.boxesNum, "boxes": []}
    for i in range(c_box.boxesNum):
        bbox_instance = c_box.boxes[i]
        bbox_dict = {
            "xmin": bbox_instance.xmin,
            "ymin": bbox_instance.ymin,
            "xmax": bbox_instance.width + bbox_instance.xmin,
            "ymax": bbox_instance.height + bbox_instance.ymin,
            "score": bbox_instance.score,
            "label": bbox_instance.label.decode("utf-8"),
        }
        box_dict["boxes"].append(bbox_dict)

    # Serialize the dictionary to JSON
    with open(output_path, "a", encoding="utf-8") as f:
        f.seek(0, os.SEEK_END)
        file_size = f.tell()
        if file_size < 500 * 1024 * 1024:
            json.dump(box_dict, f, indent=2, ensure_ascii=False)
            f.write("\n")
        else:
            print("only support 500M.", output_path, "now file size:", file_size)

def lyn_free_callback(params):
    for p in params:
        common.error_check(sdk.lyn_free(p))
    return 0

class yolov8(infer_process):
    def __init__(self, attr):
        super().__init__(attr)
        self.seg_box_apu_ptr, ret = sdk.lyn_malloc(ctypes.sizeof(LynSegBoxesInfo))
        common.error_check(ret)
        self.osd_plugin, ret = sdk.lyn_plugin_register(attr.osd_plugin)
        common.error_check(ret, "lyn_plugin_register")

        self.handle = ctypes.CDLL(attr.post_plugin)
        if not self.handle:
            raise("error: post plugin load fail.")
        self.handle.lynSegmentPostProcess.argtypes = [ctypes.POINTER(SegmentPostProcessInfo_t)]
        self.handle.lynSegmentPostProcess.restype = ctypes.c_int
        self.segmentPostFunc = self.handle.lynSegmentPostProcess

        self.handle.deleteMaskData.argtypes = [ctypes.POINTER(LynSegBoxesInfo)]
        self.handle.deleteMaskData.restype = ctypes.c_int
        self.deleteMaskDataFunc = self.handle.deleteMaskData

        if not self.segmentPostFunc or not self.deleteMaskDataFunc:
            self.handle.close()
            raise("error: post plugin function err.")

    def set_padding_data(self, video_width, video_heigh, model_width, model_height):
        
        self.ipeYolov8 = ipeParamYolov8.IpeParamYolov8(model_width, model_height)
        self.ipeYolov8.set_img_info(video_width, video_heigh,
                            sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12)
        self.ipe_bgr = ipeParamYolov8.IpeParamC2C(sdk.lyn_pixel_format_t.LYN_PIX_FMT_BGR24)
        self.ipe_nv12 = ipeParamYolov8.IpeParamC2C(sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12)

        return 0,0,0,0

    def ipe_process(self, cb_data: framepool_cb_data):
        sdk.lyn_set_current_context(self.get_ctx())
        ipe_out_data = self.ipe_output_mem_pool.pop()
        self.ipeYolov8.calc_param(self.ipe_stream, cb_data.frame.data, ipe_out_data)

        return ipe_out_data

    def before_ctx_destroy(self):
        self.handle.close()
        self.ipeYolov8.destory()
        self.ipe_bgr.destory()
        self.ipe_nv12.destory()

        common.error_check(sdk.lyn_free(self.seg_box_apu_ptr))
        ret = sdk.lyn_plugin_unregister(self.osd_plugin)
        common.error_check(ret, "lyn_plugin_unregister")

    def plugin_process(self, apu_output_data, cb_data):
        infer_host_np = np.zeros(self.apu_output_size, dtype=np.uint8)
        infer_host_ptr = sdk.lyn_numpy_to_ptr(infer_host_np)
        common.error_check(
            sdk.lyn_memcpy_async(
                self.apu_stream,
                infer_host_ptr,
                apu_output_data,
                self.apu_output_size,
                sdk.lyn_memcpy_dir_t.ServerToClient,
            )
        )
        # 等待apu处理完成
        ret = sdk.lyn_synchronize_stream(self.apu_stream)
        common.error_check(ret, "lyn_synchronize_stream")

        # ret = sdk.lyn_record_event(self.apu_stream, self.apu_event)
        # common.error_check(ret, "lyn_record_event")
        # ret = sdk.lyn_stream_wait_event(self.plugin_stream, self.apu_event)
        # common.error_check(ret, "lyn_stream_wait_event")

        pythonapi.PyCapsule_GetPointer.restype = c_void_p
        pythonapi.PyCapsule_GetPointer.argtypes = [py_object, c_char_p]
        frame_data_ptr = pythonapi.PyCapsule_GetPointer(cb_data.frame.data, None)

        seg_box_host_np = np.zeros(ctypes.sizeof(LynSegBoxesInfo), dtype=np.uint8)
        seg_box_host_ptr = sdk.lyn_numpy_to_ptr(seg_box_host_np)
        seg_box_host_ptr_c = pythonapi.PyCapsule_GetPointer(seg_box_host_ptr, None)
        infer_host_ptr_c = pythonapi.PyCapsule_GetPointer(infer_host_ptr, None)

        post_para = SegmentPostProcessInfo_t()
        post_para.model_width = self.model_width
        post_para.model_height = self.model_height
        post_para.img_width = self.codec_para.width
        post_para.img_height = self.codec_para.height
        post_para.padx = self.ipeYolov8.m_iPadX
        post_para.pady = self.ipeYolov8.m_iPadY
        post_para.fatio = self.ipeYolov8.m_fRatio
        post_para.output_tensor = infer_host_ptr_c
        post_para.boxesInfo = seg_box_host_ptr_c

        self.segmentPostFunc(post_para)

        common.error_check(
            sdk.lyn_memcpy(
                self.seg_box_apu_ptr,
                seg_box_host_ptr,
                ctypes.sizeof(LynSegBoxesInfo),
                sdk.lyn_memcpy_dir_t.ClientToServer,
            )
        )
        seg_box_apu_ptr_c = pythonapi.PyCapsule_GetPointer(self.seg_box_apu_ptr, None)
        seg_box_info = ctypes.cast(
            seg_box_host_ptr_c, ctypes.POINTER(LynSegBoxesInfo)
        ).contents

        if seg_box_info.maskDataLen > 0:
            dev_mask_data, ret = sdk.lyn_malloc(
                seg_box_info.maskDataLen
            )
            common.error_check(ret)

            addr = ctypes.addressof(ctypes.cast(seg_box_info.maskData, UINT8_PTR).contents)
            mask_data_np = np.frombuffer(
                (ctypes.c_byte * seg_box_info.maskDataLen).from_address(addr), 
                dtype=np.uint8
            )

            # buffer = ctypes.cast(seg_box_info.maskData, ctypes.POINTER(ctypes.c_uint8 * seg_box_info.maskDataLen)).contents
            # mask_data_np = np.frombuffer(
            #     buffer,
            #     dtype=np.uint8,
            # )

            mask_data_ptr_host = sdk.lyn_numpy_to_ptr(mask_data_np)

            common.error_check(
                sdk.lyn_memcpy(
                    dev_mask_data,
                    mask_data_ptr_host,
                    seg_box_info.maskDataLen,
                    sdk.lyn_memcpy_dir_t.ClientToServer,
                )
            )
            
            self.deleteMaskDataFunc(seg_box_info)

            if self.attr.show_type == 2:
                ret = sdk.lyn_stream_add_callback(
                    self.plugin_stream,
                    save_boxinfo_cb,
                    [self.attr.output_path, self.seg_box_apu_ptr, cb_data.frame.eos],
                )
            else:
                osd_para = struct.pack(
                    "1P3i1P4i",
                    seg_box_apu_ptr_c,
                    self.codec_para.width,
                    self.codec_para.height,
                    sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
                    frame_data_ptr,
                    24,
                    4,
                    2,
                    4,
                )
                common.error_check(
                    sdk.lyn_plugin_run_async(
                        self.plugin_stream, self.osd_plugin, "lynDrawSegBoxAndText", osd_para, len(osd_para)
                    )
                )

                common.error_check(sdk.lyn_synchronize_stream(self.plugin_stream))

                imgBgrSize = self.codec_para.width * self.codec_para.height * 3
                img_bgr_p, ret = sdk.lyn_malloc(imgBgrSize)
                common.error_check(ret)
                self.ipe_bgr.set_img_info(self.codec_para.width, self.codec_para.height, sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12)
                self.ipe_bgr.calc_param(self.plugin_stream, cb_data.frame.data, img_bgr_p)

                img_bgr_p_c = pythonapi.PyCapsule_GetPointer(img_bgr_p, None)
                dev_mask_data_c = pythonapi.PyCapsule_GetPointer(dev_mask_data, None)

                osd_mask_para = struct.pack(
                    "1P3i1f2P1i",
                    img_bgr_p_c,
                    self.codec_para.width,
                    self.codec_para.height,
                    3,
                    0.5,
                    seg_box_apu_ptr_c,
                    dev_mask_data_c,
                    seg_box_info.maskDataLen
                )

                common.error_check(
                    sdk.lyn_plugin_run_async(
                        self.plugin_stream, self.osd_plugin, "maskDraw", osd_mask_para, len(osd_mask_para)
                    )
                )

                self.ipe_nv12.set_img_info(self.codec_para.width, self.codec_para.height, sdk.lyn_pixel_format_t.LYN_PIX_FMT_BGR24)
                self.ipe_nv12.calc_param(self.plugin_stream, img_bgr_p, cb_data.frame.data)

                ret = sdk.lyn_stream_add_callback(
                    self.plugin_stream,
                    lyn_free_callback,
                    [img_bgr_p],
                )
                common.error_check(ret, "lyn_stream_add_callback")

            ret = sdk.lyn_stream_add_callback(
                self.plugin_stream,
                lyn_free_callback,
                [dev_mask_data],
            )
            common.error_check(ret, "lyn_stream_add_callback")

        ret = sdk.lyn_stream_add_callback(
            self.plugin_stream,
            free_to_pool_callback,
            [self.apu_output_mem_pool, apu_output_data],
        )
        common.error_check(ret, "lyn_stream_add_callback")

    def run(self, attr: infer_process_attr):
        super().run(attr)


def main(args, video_frame) -> None:
    attr = infer_process_attr()
    attr.url = args.input_path
    if args.show_type == 1:
        attr.output_path = common.append_text_to_filename(
            common.change_file_extension(args.input_path, ".264"),
            f"_device_{args.device_id}_channel_{args.channel_id}",
        )
    elif args.show_type == 2:
        attr.output_path = common.append_text_to_filename(
            common.change_file_extension(args.input_path, ".json"),
            f"_device_{args.device_id}_channel_{args.channel_id}",
        )
    if not os.path.isfile(attr.url):
        attr.output_path = attr.output_path.replace("/", "_").replace(":", "_")
    if os.path.exists(attr.output_path):
        os.remove(attr.output_path)
    attr.device_id = args.device_id
    attr.chan_id = args.channel_id
    attr.plugin_path = args.osd_plugin
    attr.model_path = args.model_path
    attr.show_type = args.show_type
    attr.video_frame = video_frame
    attr.osd_plugin = args.osd_plugin
    attr.post_plugin = args.post_plugin
    if not attr.url:
        raise ValueError("input file miss!!!")
    if attr.show_type == 1 and not attr.output_path:
        raise ValueError("unspecified output path!!!")
    yolov8_worker = yolov8(attr)
    yolov8_worker.run(cancel_flag)
    yolov8_worker.close()


def cancel_process(signum, frame):
    global cancel_flag
    cancel_flag.value = True


if __name__ == "__main__":
    argparser = argparse.ArgumentParser(description="python3 main.py ")
    argparser.add_argument("-m", "--model_path", type=str, help="model file path")
    argparser.add_argument("-i", "--input_path", help="input video path or stream url")
    argparser.add_argument("-p", "--post_plugin", type=str, help="plugin path")
    argparser.add_argument("-o", "--osd_plugin", type=str, help="plugin path")
    argparser.add_argument(
        "-d",
        "--device_id",
        type=str,
        default="0",
        help='Comma-separated list of device ids, e.g., "0,1"',
    )
    argparser.add_argument("-s", "--show_type", type=int, default=0, help="show type")
    argparser.add_argument(
        "-c", "--channel_num", type=int, default=1, help="channel_num"
    )

    signal.signal(signal.SIGINT, cancel_process)
    args = argparser.parse_args()
    args.config_file = None
    channel_infos = get_channel_infos(args)
    start_process(main, channel_infos, cancel_flag)
